package com.tripcube.ai.service.impl;

import com.aliyun.bailian20231229.Client;
import com.aliyun.bailian20231229.models.*;
import com.aliyun.teautil.models.RuntimeOptions;
import com.tripcube.ai.model.dto.CreateKnowledgeBaseDTO;
import com.tripcube.ai.service.KnowledgeBaseService;
import com.tripcube.common.core.exception.ServiceException;
import lombok.RequiredArgsConstructor;
import lombok.extern.slf4j.Slf4j;
import org.springframework.beans.factory.annotation.Value;
import org.springframework.stereotype.Service;

import java.util.HashMap;
import java.util.List;
import java.util.Map;
import java.util.Objects;

/**
 * 知识库管理服务实现类
 *
 * @author tripcube
 */
@Service
@RequiredArgsConstructor
@Slf4j
public class KnowledgeBaseServiceImpl implements KnowledgeBaseService {

    private final Client bailianClient;

    @Value("${alibaba.cloud.workspace-id}")
    private String workspaceId;

    /**
     * 创建知识库
     */
    @Override
    public Object createKnowledgeBase(CreateKnowledgeBaseDTO createDTO) {
        try {
            // 构建创建知识库请求
            CreateIndexRequest request = new CreateIndexRequest()
                    .setName(createDTO.getName())
                    .setDescription(createDTO.getDescription())
                    .setStructureType(createDTO.getStructureType())
                    .setSinkType(createDTO.getSinkType())
                    .setEmbeddingModelName(createDTO.getEmbeddingModelName());

            if(Objects.equals(createDTO.getSourceType(), "DATA_CENTER_FILE")){
                request.setDocumentIds(createDTO.getDocumentIds());
            } else if (Objects.equals(createDTO.getSourceType(), "DATA_CENTER_CATEGORY")) {
                request.setCategoryIds(createDTO.getCategoryIds());
            }
            RuntimeOptions runtime = new RuntimeOptions();
            Map<String, String> headers = new HashMap<>();

            // 创建知识库
            CreateIndexResponseBody body = bailianClient.createIndexWithOptions(
                    workspaceId, request, headers, runtime).body;

            if(body.data == null){
                log.error("创建知识库失败: {}", body.getMessage());
                throw new ServiceException("创建知识库失败: " + body.getMessage());
            }

            String indexId = body.getData().getId();

            // 提交知识库索引任务
            SubmitIndexJob(indexId);

            return body.data;

        } catch (Exception e) {
            log.error("创建知识库失败", e);
            throw new ServiceException("创建知识库失败: " + e.getMessage());
        }
    }

    /**
     * 查询知识库列表
     */
    @Override
    public Object selectKnowledgeBaseList(String indexName,String pageNumber,String pageSize) {
        try {
            ListIndicesRequest request = new ListIndicesRequest();
            request.setIndexName(indexName);
            request.setPageNumber(pageNumber);
            request.setPageSize(pageSize);
            RuntimeOptions runtime = new RuntimeOptions();

            // 调用阿里百炼API查询知识库列表
            ListIndicesResponseBody body = bailianClient.listIndicesWithOptions(
                    workspaceId, request, new HashMap<>(), runtime).body;

            if(body.data != null){
                return body.data;
            }else {
                log.error("查询知识库列表失败: {}", body.getMessage());
                throw new ServiceException("查询知识库列表失败: " + body.getMessage());
            }

        } catch (Exception e) {
            log.error("查询知识库列表失败", e);
            throw new ServiceException("查询知识库列表失败: " + e.getMessage());
        }
    }

    /**
     * 删除知识库
     */
    @Override
    public int deleteKnowledgeBase(String knowledgeBaseId) {
        try {
            DeleteIndexRequest request = new DeleteIndexRequest()
                    .setIndexId(knowledgeBaseId);
            RuntimeOptions runtime = new RuntimeOptions();

            DeleteIndexResponse response = bailianClient.deleteIndexWithOptions(
                    workspaceId, request, new HashMap<>(), runtime);

            if (response.getBody() != null && "Success".equals(response.getBody().getCode())) {
                log.info("成功删除知识库，ID: {}", knowledgeBaseId);
                return 1;
            }

            return 0;

        } catch (Exception e) {
            log.error("删除知识库失败，ID: {}", knowledgeBaseId, e);
            throw new ServiceException("删除知识库失败: " + e.getMessage());
        }
    }

    /**
     * 批量删除知识库
     */
    @Override
    public int deleteKnowledgeBaseBatch(String[] knowledgeBaseIds) {
        int count = 0;
        for (String id : knowledgeBaseIds) {
            count += deleteKnowledgeBase(id);
        }
        return count;
    }


    /**
     * 查询索引下的文档列表
     */
    @Override
    public Object listIndexDocuments(String indexId, String documentStatus, String documentName, Integer pageNumber, Integer pageSize) throws Exception {
        ListIndexDocumentsRequest request = new ListIndexDocumentsRequest();
        request.setIndexId(indexId);
        if (documentStatus != null) {
            request.setDocumentStatus(documentStatus);
        }
        if (documentName != null) {
            request.setDocumentName(documentName);
        }
        request.setPageNumber(pageNumber);
        request.setPageSize(pageSize);

        request.setEnableNameLike(String.valueOf(true));

        RuntimeOptions runtime = new RuntimeOptions();
        Map<String, String> headers = new HashMap<>();
        return bailianClient.listIndexDocumentsWithOptions(workspaceId, request, headers, runtime).body.data;
    }

    /**
     * 删除知识库索引文档
     */
    @Override
    public Object deleteIndexDocument(String indexId, List<String> documentIds) throws Exception {
        DeleteIndexDocumentRequest request = new DeleteIndexDocumentRequest();
        request.setIndexId(indexId);
        request.setDocumentIds(documentIds);

        RuntimeOptions runtime = new RuntimeOptions();
        Map<String, String> headers = new HashMap<>();
        return bailianClient.deleteIndexDocumentWithOptions(workspaceId, request, headers, runtime);
    }


    /**
     * 提交索引任务
     */
    @Override
    public Object SubmitIndexJob(String indexId) {
        SubmitIndexJobRequest submitIndexJobRequest = new SubmitIndexJobRequest();
        submitIndexJobRequest.setIndexId(indexId);
        RuntimeOptions runtime = new com.aliyun.teautil.models.RuntimeOptions();
        try {
            SubmitIndexJobResponseBody body = bailianClient.submitIndexJobWithOptions(workspaceId, submitIndexJobRequest, new HashMap<>(), runtime).body;
            if(body.data != null){
                return body.data;
            }else{
                log.error(body.getMessage());
                throw new ServiceException("提交索引任务失败: " + body.getMessage());
            }
        } catch (Exception e) {
            log.error("提交索引任务失败，ID: {}", indexId, e);
            throw new ServiceException("提交索引任务失败: " + e.getMessage());
        }
    }

    /**
     * 查询索引任务状态
     */
    @Override
    public Object getIndexJobStatus(String indexId, String jobId) {
        GetIndexJobStatusRequest getIndexJobStatusRequest = new GetIndexJobStatusRequest();
        getIndexJobStatusRequest.setIndexId(indexId);
        getIndexJobStatusRequest.setJobId(jobId);
        Map<String, String> headers = new HashMap<>();
        RuntimeOptions runtime = new RuntimeOptions();
        try {
            GetIndexJobStatusResponseBody body = bailianClient.getIndexJobStatusWithOptions(workspaceId, getIndexJobStatusRequest, headers, runtime).body;
            if(body.data != null){
                return body.data;
            }else{
                log.error(body.getMessage());
                throw new ServiceException("查询索引任务状态失败: " + body.getMessage());
            }
        } catch (Exception e) {
            log.error("查询索引任务状态失败，ID: {}", indexId, e);
            throw new ServiceException("查询索引任务状态失败: " + e.getMessage());
        }
    }

     /**
     * 提交索引追加任务
     */
    @Override
    public Object submitIndexAddDocumentsJob(String indexId, String sourceType, List<String> categoryIds, List<String> documentIds) {
        SubmitIndexAddDocumentsJobRequest submitIndexAddDocumentsJobRequest = new SubmitIndexAddDocumentsJobRequest();
        submitIndexAddDocumentsJobRequest.setIndexId(indexId);
        submitIndexAddDocumentsJobRequest.setSourceType(sourceType);
        if(sourceType.equals("DATA_CENTER_CATEGORY")){
            submitIndexAddDocumentsJobRequest.setCategoryIds(categoryIds);
        }else if (sourceType.equals("DATA_CENTER_File")){
            submitIndexAddDocumentsJobRequest.setDocumentIds(documentIds);
        }
        Map<String, String> headers = new HashMap<>();
        RuntimeOptions runtime = new RuntimeOptions();
        try {
            SubmitIndexAddDocumentsJobResponseBody body = bailianClient.submitIndexAddDocumentsJobWithOptions(workspaceId, submitIndexAddDocumentsJobRequest, headers, runtime).body;
            if(body.data != null){
                return body.data;
            }else{
                log.error(body.getMessage());
                throw new ServiceException("提交索引追加任务失败: " + body.getMessage());
            }
        } catch (Exception e) {
            log.error("提交索引追加任务失败，ID: {}", indexId, e);
            throw new ServiceException("提交索引追加任务失败: " + e.getMessage());
        }
    }

    /**
     * 查询索引下的分片列表
     */
    @Override
    public Object listChunks(String indexId, String fileId, Integer pageNum, Integer pageSize, List<String> fields) throws Exception {
        ListChunksRequest request = new ListChunksRequest();
        request.setIndexId(indexId);
        if (fileId != null) {
            request.setFileId(fileId);
        }
        request.setPageNum(pageNum);
        request.setPageSize(pageSize);
        if (fields != null && !fields.isEmpty()) {
            request.setFields(fields);
        }

        RuntimeOptions runtime = new RuntimeOptions();
        Map<String, String> headers = new HashMap<>();
        ListChunksResponseBody body = bailianClient.listChunksWithOptions(workspaceId, request, headers, runtime).body;
        if(body.data != null){
            return body.data;
        }else{
            log.error(body.getMessage());
            throw new ServiceException("查询索引下的分片列表失败: " + body.getMessage());
        }
    }

    /**
     * 更新索引下的分片
     */
    @Override
    public Object updateChunk(String indexId, String dataId, String chunkId, Boolean isDisplayedChunkContent, String content, String title) throws Exception {
        UpdateChunkRequest request = new UpdateChunkRequest();
        request.setPipelineId(indexId);
        request.setDataId(dataId);
        request.setChunkId(chunkId);
        request.setIsDisplayedChunkContent(isDisplayedChunkContent);
        request.setContent(content);
        request.setTitle(title);

        RuntimeOptions runtime = new RuntimeOptions();
        Map<String, String> headers = new HashMap<>();
        UpdateChunkResponseBody body = bailianClient.updateChunkWithOptions(workspaceId, request, headers, runtime).body;
        if(body.data != null){
            return body.data;
        }else{
            log.error(body.getMessage());
            throw new ServiceException("更新索引下的分片失败: " + body.getMessage());
        }
    }

    /**
     * 删除索引下的分片
     */
    @Override
    public Object deleteChunks(String indexId, List<String> chunkIds) throws Exception {
        DeleteChunkRequest request = new DeleteChunkRequest();
        request.setPipelineId(indexId);
        request.setChunkIds(chunkIds);

        RuntimeOptions runtime = new RuntimeOptions();
        Map<String, String> headers = new HashMap<>();
        DeleteChunkResponseBody body = bailianClient.deleteChunkWithOptions(workspaceId, request, headers, runtime).body;
        if(body.data != null){
            return body.data;
        }else{
            log.error(body.getMessage());
            throw new ServiceException("删除索引下的分片失败: " + body.getMessage());
        }
    }
}